# import torch
# from Model.model2 import CRNN
# import struct
#
# # model_path = './data/crnn.pth'
#
# model = CRNN()
# stata_dict=torch.load("ch_rec_server_crnn_res34.pth")['state_dict']
# model.load_state_dict(stata_dict)
# if torch.cuda.is_available():
#     model = model.cuda()
# # print('loading pretrained model from %s' % model_path)
# # model.load_state_dict(torch.load(model_path))
#
# image = torch.ones(1, 3, 32, 120)
# if torch.cuda.is_available():
#     image = image.cuda()
#
# model.eval()
# print(model)
# print('image shape ', image.shape)
# preds = model(image)
#
# f = open("crnn.wts", 'w')
# f.write("{}\n".format(len(model.state_dict().keys())))
# for k,v in model.state_dict().items():
#     print('key: ', k)
#     print('value: ', v.shape)
#     vr = v.reshape(-1).cpu().numpy()
#     f.write("{} {}".format(k, len(vr)))
#     for vv in vr:
#         f.write(" ")
#         f.write(struct.pack(">f", float(vv)).hex())
#     f.write("\n")

# -*- coding: utf-8 -*-
"""
Created on Sat May 29 19:12:58 2021
@author: Administrator
"""
import xml.etree.ElementTree as ET
import os
from PIL import Image
import numpy as np
import torch

from PIL import Image
import torch
import torchvision
import struct

path = 'Wts'
Path = os.path.join(path, "wts")
if not os.path.isdir(Path):
    os.makedirs(Path)


def getweights(model_path):
    state_dict = torch.load(model_path, map_location=lambda storage, loc: storage)['state_dict']
    # print(state_dict)
    keys = [v for key, v in enumerate(state_dict)]
    # print(keys)
    with open(os.path.join(Path, "network.txt"), 'w') as fw:
        for key in keys:
            print("~~~~~~~~~~~ ", key)
            ts = state_dict[key]
            shape = ts.shape
            size = shape
            allsize = 1
            fw.write(key + " ")
            for idx in range(len(size)):
                allsize *= size[idx]
                fw.write(str(size[idx]) + " ")
            fw.write('\n')
            ts = ts.reshape(allsize)
            with open(Path + '/' + key + '.wgt', 'wb') as f:
                a = struct.pack('i', allsize)
                f.write(a)
                for i in range(allsize):
                    a = struct.pack('f', ts[i])  # .hex()
                    f.write(a)


if __name__ == '__main__':
    from model import DBNet
    model = DBNet()
    import torch

    state_dict = torch.load("ch_det_mobile_db_mbv3.pth")
    model.load_state_dict(state_dict['state_dict'])
    x = torch.ones([1, 3, 640, 640])
    # import torch.onnx
    # torch.onnx.export(model,x,'dbnet.onnx',verbose=True,training=2)
    # import netron
    # netron.start("dbnet.onnx")


    torch.save(model, path + 'dbnet.pth')
    getweights("ch_det_mobile_db_mbv3.pth")
    # getweights(path + "dbnet.pth")
    # model = torchvision.models.resnet50()
    # model.eval()
    # torch.save(model.state_dict(),r"H:\myGitHub\tensorrtF\model\resnet50\res50.pth")
    # a = torch.randn(1,3,256,256).type(torch.float32)
    # torch.onnx.export(model, a,r"H:\myGitHub\tensorrtF\model\resnet50\res50.onnx",training=2 )state_dict